/*
* Copyright (c) 2025 Shenzhen Kaihong Digital Industry Development Co., Ltd.
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

use proc_macro::TokenStream;
use syn::{Data, DeriveInput};

#[proc_macro_derive(Encode, attributes(encode))]
pub fn derive_encode(input: TokenStream) -> TokenStream {
    let input = syn::parse_macro_input!(input as DeriveInput);
    match &input.data {
        Data::Struct(data_struct) => {
            let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
            let ty = &input.ident;

            let encode_unchecked_stmts = data_struct
                .fields
                .iter()
                .filter_map(|field| {
                    if let Some(ident) = &field.ident {
                        Some(quote::quote! {
                            unsafe { self.#ident.encode_unchecked(encoder, buf) };
                        })
                    } else {
                        None
                    }
                });
            let encode_stmts = data_struct
                .fields
                .iter()
                .filter_map(|field| {
                    if let Some(ident) = &field.ident {
                        Some(quote::quote! {
                            self.#ident.encode(encoder, buf)?;
                        })
                    } else {
                        None
                    }
                });

            quote::quote! {
                impl #impl_generics ohmw_core::bincode::Encode for #ty #ty_generics #where_clause {
                    unsafe fn encode_unchecked<E: ohmw_core::bincode::Encoder>(&self, encoder: &mut E, buf: &mut [u8]) {
                        #(#encode_unchecked_stmts)*
                    }

                    fn encode<E: ohmw_core::bincode::Encoder>(&self, encoder: &mut E, buf: &mut Vec<u8>) -> Result<(), E::Error> {
                        #(#encode_stmts)*
                        Ok(())
                    }
                }
            }.into()
        },
        _ => panic!("Only support struct"),
    }
}

#[proc_macro_derive(Decode, attributes(decode))]
pub fn derive_decode(input: TokenStream) -> TokenStream {
    let input = syn::parse_macro_input!(input as DeriveInput);
    match &input.data {
        Data::Struct(data_struct) => {
            let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
            let ty = &input.ident;

            let decode_unchecked_stmts = data_struct
                .fields
                .iter()
                .filter_map(|field| {
                    if let Some(ident) = &field.ident {
                        let ty = &field.ty;
                        Some(quote::quote! {
                            let #ident = unsafe { <#ty as ohmw_core::bincode::Decode>::decode_unchecked(decoder, buf) };
                        })
                    } else {
                        None
                    }
                });
            let decode_stmts = data_struct
                .fields
                .iter()
                .filter_map(|field| {
                    if let Some(ident) = &field.ident {
                        let ty = &field.ty;
                        Some(quote::quote! {
                            let #ident = <#ty as ohmw_core::bincode::Decode>::decode(decoder, buf)?;
                        })
                    } else {
                        None
                    }
                });
            let fields_with_comma = data_struct
                .fields
                .iter()
                .filter_map(|field| {
                    if let Some(ident) = &field.ident {
                        Some(quote::quote! {
                            #ident,
                        })
                    } else {
                        None
                    }
                })
                .collect::<Vec<_>>();

            quote::quote! {
                impl #impl_generics ohmw_core::bincode::Decode for #ty #ty_generics #where_clause {
                    unsafe fn decode_unchecked<D: ohmw_core::bincode::Decoder>(decoder: &mut D, buf: &[u8]) -> Self {
                        #(#decode_unchecked_stmts)*
                        Self {
                            #(#fields_with_comma)*
                        }
                    }

                    fn decode<D: ohmw_core::bincode::Decoder>(decoder: &mut D, buf: &[u8]) -> Result<Self, D::Error> {
                        #(#decode_stmts)*
                        Ok(Self {
                            #(#fields_with_comma)*
                        })
                    }
                }
            }.into()
        },
        _ => panic!("Only support struct"),
    }
}

#[proc_macro_derive(BorrowDecode, attributes(borrow_decode))]
pub fn derive_borrow_decode(input: TokenStream) -> TokenStream {
    let input = syn::parse_macro_input!(input as DeriveInput);
    match &input.data {
        Data::Struct(data_struct) => todo!(),
        _ => panic!("Only support struct"),
    }
}
